Skip to content

feat(search): add multi-query batch search support - #2607

Open
LHT129 wants to merge 1 commit into
antgroup:mainfrom
LHT129:batch_query_review_reset_1685
Open

feat(search): add multi-query batch search support#2607
LHT129 wants to merge 1 commit into
antgroup:mainfrom
LHT129:batch_query_review_reset_1685

Conversation

@LHT129

@LHT129 LHT129 commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Replacement for #1685, opened from the same current commit to start a clean review. This change adds multi-query batch KNN support for HGraph and IVF, updates API semantics and regression coverage, while retaining single-query behavior where batch result shapes are not representable. Closes #1684.

Copilot AI lite review requested due to automatic review settings August 3, 2026 09:46
@vsag-bot

vsag-bot commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

/label status/waiting-for-review
/waiting-on reviewer
/request-review @jiaweizone
/request-review @wxyucs
/request-review @inabao

@vsag-bot

vsag-bot commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Automated pull request review completed.

Review effort: high (1428 changed lines across 19 files).

Submitted 4 inline comments.
Review: #2607 (review)

@mergify

mergify Bot commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Merge Protections

🟢 All 3 merge protections satisfied — ready to merge.

Show 3 satisfied protections

🟢 Require kind label

  • label~=^kind/

🟢 Require version label

  • label~=^version/

🟢 Require linked issue for feature/bug PRs

  • body~=(?im)(?:^|[\s\-\*])(?:close[sd]?|fix(?:e[sd])?|resolve[sd]?)\s*:?\s+(?:#\d+|[\w.\-]+/[\w.\-]+#\d+|https?://github\.com/[\w.\-]+/[\w.\-]+/issues/\d+)

Comment thread src/algorithm/hgraph/hgraph_search.cpp Outdated
Comment thread src/algorithm/ivf/ivf.cpp Outdated
Comment thread src/algorithm/hgraph/hgraph_search.cpp
Comment thread src/algorithm/hgraph/hgraph_search.cpp
Comment thread src/impl/label_table/label_table.h Outdated
Comment thread src/algorithm/hgraph/hgraph_search.cpp Outdated

@LHT129 LHT129 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall this is a well-structured PR that adds multi-query batch KNN search support for HGraph and IVF. The code quality is high with thorough overflow guards, sentinel-based padding, and clear documentation updates.

Summary of findings:

  1. [suggestion] Range search statistics regression: The extracted search_range_with_request method uses ctx.stats->Dump() instead of mci_result.MakeStatistics(stats).Dump(), losing the route field (brute_force/mci/graph) from range search statistics output.

  2. [suggestion] IVF batch path Dataset::Make() per iteration: The temporary dataset allocation inside the per-query loop could be hoisted out for a minor performance improvement.

  3. [note] HasActiveLabel naming: The function is hardcoded for label -1 but has a general-purpose signature. Consider renaming to HasActivePaddingLabel().

  4. [note] last_result_inner_ids naming: The variable name is slightly misleading since it only captures reasoning-related inner IDs (single-query only). Consider renaming to reasoning_inner_ids.

Positive observations:

  • Comprehensive overflow guards for query_count * k and byte-level allocations
  • Sentinel pre-fill with ids = -1 and dists = +inf is well-designed
  • Clean extraction of range search into search_range_with_request
  • Good test coverage including empty index, batch KNN, batch range rejection, and IVF bucket routing
  • Proper rejection of reasoning with batch queries
  • label_table tracking of -1 labels is correctly maintained across all mutation paths (Insert, Remove, Merge, UpdateLabel, ShrinkToFit, Deserialize)

@LHT129 LHT129 added kind/feature Brand-new functionality or capabilities 引入全新的功能、新特性或新能力 version/1.1 labels Aug 3, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds multi-query (batched) KNN search support to VSAG’s core indexes (notably HGraph and IVF) by allowing DatasetPtr queries with NumElements > 1, defining/clarifying result layout semantics, and extending regression coverage and API documentation accordingly.

Changes:

  • Implement batched KNN execution paths for HGraph and IVF (with explicit rejection of batched range search).
  • Standardize batch result layout to row-major query_count x dim with sentinel padding (id = -1), and enforce the “no external label -1” constraint for unambiguous padding.
  • Add/extend functional tests and update public API docs to describe single-query vs batched semantics.

Reviewed changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
tests/test_ivf.cpp Adds multi-query KNN tests (including empty-index behavior) and batch routing behavior checks for IVF.
tests/test_hgraph.cpp Adds/extends multi-query KNN tests and asserts multi-query range search is rejected for HGraph.
src/utils/timer.h Adds Timer::Reset() declaration (used for per-query timeout tracking in batch search).
src/utils/timer.cpp Implements Timer::Reset().
src/index/index_impl.h Adjusts empty-index short-circuit behavior to better support batch-query semantics.
src/impl/label_table/label_table.h Tracks active external label -1 usage to gate batch KNN padding semantics.
src/impl/label_table/label_table.cpp Wires allocator/maintenance of the active-padding-label tracker and rebuilds it on deserialize/merge.
src/algorithm/ivf/ivf.cpp Implements IVF batch KNN behavior inside SearchWithRequest, including padding and overflow guards.
src/algorithm/hgraph/hgraph.h Extends get_data to support indexed query access and adds offset overflow guards.
src/algorithm/hgraph/hgraph_serialize.cpp Rebuilds active padding label tracking after legacy label-table deserialization paths.
src/algorithm/hgraph/hgraph_search.cpp Implements HGraph batch KNN in SearchWithRequest, factors out range-search path, and adds padding/overflow handling.
include/vsag/search_request.h Updates SearchRequest::query_ docs to describe single vs batched semantics and constraints.
include/vsag/index.h Updates Index::SearchWithRequest result-shape documentation for single vs batched behaviors.
docs/docs/zh/src/api/search.md Documents batched KNN availability/limitations in the Chinese API docs.
docs/docs/zh/src/api/index_class.md Updates Chinese index API docs for batched KNN result reading and constraints.
docs/docs/zh/src/api/dataset.md Updates Chinese dataset docs to explain batched result matrix layout and padding.
docs/docs/en/src/api/search.md Documents batched KNN availability/limitations and clarifies IVF routing-only mode wording.
docs/docs/en/src/api/index_class.md Updates English index API docs for batched KNN result reading and constraints.
docs/docs/en/src/api/dataset.md Updates English dataset docs to explain batched result matrix layout and padding.
Suppressed comments (1)

include/vsag/search_request.h:51

  • Same indentation issue continues in the remainder of this bullet; keeping the alignment consistent avoids broken formatting in generated API docs.
      *            fewer neighbors than the returned Dim are padded with sentinel entries
      *            (id = -1, distance = +infinity). Batch KNN rejects an index containing external
      *            label -1 to keep this padding unambiguous.
     *          - Batched RANGE_SEARCH is not supported; implementations MUST reject
     *            NumElements > 1 for range mode.

Comment thread src/utils/timer.h
Comment thread include/vsag/index.h
Comment thread include/vsag/search_request.h Outdated

@vsag-bot vsag-bot left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated inline review completed.

Review effort: high (1150 changed lines across 19 files).
Submitted 2 inline comments.

Reviewed commit 500b5ca.

Comment thread include/vsag/search_request.h
Comment thread src/algorithm/hgraph/hgraph_search.cpp Outdated
Comment thread src/algorithm/hgraph/hgraph.h
Comment thread src/algorithm/ivf/ivf.cpp Outdated
Comment thread src/algorithm/hgraph/hgraph_search.cpp

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated 2 comments.

Suppressed comments (9)

include/vsag/index.h:340

  • The new batch contract is only documented for SearchWithRequest, but HGraph and IVF's public KnnSearch overloads delegate here and now accept the same multi-query inputs. The overload documentation above still promises num_elements = 1 and a num_elements * k layout, so callers of the primary KNN APIs receive an incorrect contract; update those overloads too.
      *                - batched KNN requests, when supported by the implementation:
      *                  num_elements = query->GetNumElements(),
      *                  dim = implementation-defined returned row width. HGraph clamps it to
      *                  min(request.topk_, GetNumElements()), while IVF preserves
      *                  request.topk_. Callers MUST read `dim` from the returned dataset.
      *                  ids/distances are stored row-major with length (num_elements * dim).
      *                  Queries that yielded fewer than dim neighbors are padded with
       *                  sentinel entries (id = -1, distance = +infinity). Batch KNN rejects

src/algorithm/hgraph/hgraph_search.cpp:116

  • This iterator-only single-query check is after the empty-index early return at line 92. IndexImpl deliberately forwards multi-query requests for the new batch handling, so a multi-query iterator call on an empty HGraph returns a single empty dataset instead of rejecting the unsupported iterator shape. Move the check before the early return.
    CHECK_ARGUMENT(query->GetNumElements() == 1,
                   "iterator-based KnnSearch only supports single query (NumElements=1)");

src/algorithm/hgraph/hgraph_search.cpp:936

  • This RaBitQ rerank call also omits request.threshold_. As in the regular rerank branch, the post-rerank filter cannot recover eligible candidates discarded when the unfiltered top-k heap was formed, so thresholded batch KNN can return incomplete results. Pass the request threshold to reorder.
        } else if (mci_result.route != "mci" && !brute_force_used && search_param.enable_reorder &&
                   params.rabitq_one_bit_search) {
            this->reorder(raw_query, this->basic_flatten_codes_, search_result, k, nullptr, ctx);

src/algorithm/hgraph/hgraph_search.cpp:1026

  • Unlike the single-query branch, which serializes mci_result.MakeStatistics(stats), the batch path always serializes only stats.Dump(). MCI-enabled batch searches therefore omit the mci_hybrid_*, seed-count, and raw-CSR diagnostics from the returned dataset, making the result statistics incomplete. Preserve or explicitly aggregate the per-query MCI metadata for batch results.
    if (query_count > 1) {
        dataset_results->Statistics(stats.Dump());
    }

src/algorithm/hgraph/hgraph_serialize.cpp:320

  • The new active-padding rebuild is present for the modern label-info paths, but deserialize_basic_info_v0_14 still reads label_table_ directly at line 170 without rebuilding it. A legacy index containing an active external label -1 will therefore bypass the new batch-safety check and can return an ambiguous -1 result; rebuild the tracking set in that legacy path too.
        this->label_table_->RebuildActivePaddingLabelIds();

docs/docs/en/src/api/search.md:106

  • This constraint is now inconsistent with the IVF implementation: batch KNN accepts one bucket list per query and validates the outer vector against the query count. Document the batch form here; otherwise callers will be told to use a shape that the new implementation deliberately supports.
- Currently only single-query is supported; the outer vector must contain exactly one entry.

docs/docs/zh/src/api/search.md:100

  • 此约束已与 IVF 实现不一致:批量 KNN 支持每个查询一个桶列表,并会校验外层向量与查询数一致。这里应记录批量形式,否则文档会要求调用方使用新实现不支持的形状。
- 当前仅支持单查询;外层向量必须恰好包含一个条目。

src/algorithm/hgraph/hgraph_search.cpp:857

  • Routing uses this separate ep_search_param, but it never receives base_search_param.time_cost; the searcher only checks InnerSearchParam::time_cost when enforcing timeout_ms. Consequently hierarchical routing in this batch path is not timeout-bounded. Propagate the request timer to the routing parameter (and reset it at the start of each query if the timeout is intended to be per-query).
    InnerSearchParam ep_search_param;
    ep_search_param.ep = this->entry_point_id_;
    ep_search_param.topk = 1;
    ep_search_param.ef = 1;
    ep_search_param.is_inner_id_allowed = nullptr;

src/algorithm/ivf/ivf.cpp:2145

  • This batch branch invokes this->SearchWithRequest once per query, and each invocation creates a separate SearchStatistics; the outer stats object used for the final dataset is never updated. The returned batch result therefore reports zero distance evaluations and misses subquery timeouts even though the searches ran. Aggregate the subrequest statistics or refactor the batch path to share the query context.
        CHECK_ARGUMENT(request.expected_labels_.empty(),
                       "IVF batch search does not support expected labels");
        CHECK_ARGUMENT(request.topk_ > 0, "topk must be greater than 0");
        CHECK_ARGUMENT(!this->label_table_->HasActivePaddingLabel(),
                       "batch KNN does not support an index containing external label -1");

Comment thread src/algorithm/hgraph/hgraph_search.cpp
Comment thread src/algorithm/hgraph/hgraph_search.cpp Outdated
if (visited_list != nullptr) {
pool->ReturnOne(visited_list);
}
}

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] When element_count == 0 and query_count > 1 (line 835-841), the returned dataset has Dim(0). This is inconsistent with the normal batch layout where Dim is k. A caller that unconditionally reads dim = result->GetDim() and indexes with q_idx * dim + i would get 0 here, which differs from the documented rectangular query_count x k layout. Consider setting Dim(k) here for consistency with the non-empty batch path, or explicitly documenting this edge case in the API contract.

The same issue applies to the k == 0 early return at line 856-862.

visited_list.reset();
}
}
FilterPtr ft = this->create_search_filter(request.filter_, params.use_extra_info_filter);

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] In search_range_with_request, when brute_force_threshold triggers the brute-force path (line 727-732), the mci_result statistics object is left with its default-constructed state (route = "", valid_ratio = 0). The final statistics at line 779 use mci_result.MakeStatistics(*ctx.stats).Dump(), which will report an empty route string for this path. The old inline code set mci_result.route = "brute_force" before calling brute_force_search so that statistics correctly reflected the search path taken.

The KNN batch path correctly sets mci_result.route = "brute_force" at line 1161, but search_range_with_request at line 732 omits this assignment.

Comment thread tests/test_ivf.cpp
check_bucket_result(batch_result.value(), 3, scan_buckets_count, buckets_count);
}

SECTION("batch routing ignores search-only options") {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] The test "batch routing ignores search-only options" at line 2377 uses RANGE_SEARCH mode with 2 queries and disable_bucket_scan params. This test passes because the bucket routing path (which handles disable_bucket_scan) returns early before reaching the range single-query validation. While this is correct behavior (bucket routing is a special mode that bypasses normal search), it may be worth adding a comment or making the test intent clearer — a reader might wonder why a 2-query RANGE_SEARCH succeeds when the API documents that range search only supports single queries.

@LHT129 LHT129 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] Overall, this PR is well-structured with solid engineering: the overflow guards, sentinel pre-fill, padding label tracking, and per-query entry point search are all correct and thorough. The existing 30+ comments already cover the critical issues (RAII visited list guards, IVF batch search implementation gaps, missing use_custom_distance guards). I added 3 additional notes:

  1. hgraph_search.cpp:835Dim(0) is returned when element_count == 0 && query_count > 1, while the single-query path returns Dim(element_count). Consider returning Dim(element_count) consistently in both paths.
  2. hgraph_search.cpp:732search_range_with_request does not set mci_result.route = "brute_force" when falling through to brute force, unlike the KNN path.
  3. test_ivf.cpp:2377 — The "batch routing ignores search-only options" test uses RANGE_SEARCH with 2 queries; consider adding a comment clarifying that bucket routing bypasses the range single-query validation.

The core batch KNN implementation is solid. The main areas to address are the existing critical comments (visited list RAII, IVF batch search, custom distance guards).

Comment thread include/vsag/search_request.h Outdated
/**
* @brief Pre-selected bucket IDs for bypassing IVF bucket routing (ClassifyDatasForSearch)
* @details The outer vector contains one entry per query vector.
* @details Currently only single-query is supported; outer vector must contain exactly one entry.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] The bucket_ids_ documentation says "Currently only single-query is supported; outer vector must contain exactly one entry", but the validation in index_impl.h (line 508) now allows bucket_ids_.size() != 1 for IVF indexes. The IVF batch path in ivf.cpp (lines 1987-2000) also handles per-query bucket_ids correctly for multi-query.

The docstring should be updated to reflect that IVF now supports multi-query bucket_ids_ in batch KNN mode.

Copilot AI review requested due to automatic review settings August 24, 2026 06:10
@LHT129
LHT129 force-pushed the batch_query_review_reset_1685 branch from 5357338 to 36a5e49 Compare August 24, 2026 06:10

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated 9 comments.

Suppressed comments (3)

include/vsag/index.h:337

  • This new batched-result contract is also reached by the non-iterator Index::KnnSearch overloads for HGraph and IVF (the added tests call those overloads with multi-element queries), but the public overload documentation above still states num_elements == 1. Update the direct KNN overload descriptions too, otherwise callers of the newly supported API are given the opposite shape contract.
      *                - batched KNN requests, when supported by the implementation:
      *                  num_elements = query->GetNumElements(),
      *                  dim = implementation-defined returned row width. HGraph clamps it to
      *                  min(request.topk_, GetNumElements()), while IVF preserves
      *                  request.topk_. Callers MUST read `dim` from the returned dataset.

src/algorithm/hgraph/hgraph_search.cpp:123

  • Because the empty-index return above precedes this new validation, an iterator call with a multi-query dataset on an empty HGraph returns a successful empty dataset instead of rejecting the unsupported batch shape. This makes the single-query-only iterator contract depend on whether the index has data; validate NumElements() before the empty-index fast path.
    // Iterator state is maintained per query, so this overload remains single-query only.
    CHECK_ARGUMENT(query->GetNumElements() == 1,
                   "iterator-based KnnSearch only supports single query (NumElements=1)");

src/impl/label_table/label_table.cpp:188

  • new_label == -1 is registered as an active padding label without checking whether source id i is removed. If a source index contains external label -1 and that vector was MARK_REMOVEd, merging it still makes the destination report an active padding label, causing all subsequent batch KNN requests to be rejected (and the source deletion state is not otherwise transferred). Carry the source keep/removal state through the merge and only register retained active labels.
            if (new_label == -1) {
                std::scoped_lock wlock(delete_ids_mutex_);
                active_padding_label_ids_.insert(new_inner_id);

For KNN, `GetNumElements()` is `1` and the ids/distances arrays have length `k`. For range search,
the number of matches is reported through the result's dimension. See
[k-Nearest Neighbor Search](../guide/knn_search.md).
For single-query KNN, `GetNumElements()` is `1` and the ids/distances arrays have length `k`. HGraph
**Constraints:**
- Batch IVF search supports KNN only; custom query distance and reasoning labels are unsupported.
- A non-empty outer vector must contain exactly one non-empty entry per query vector.
- Currently only single-query is supported; the outer vector must contain exactly one entry.

对 KNN,`GetNumElements()` 为 `1`,ids/distances 数组长度为 `k`。对范围搜索,命中数通过结果的维度报告。
见 [k-近邻搜索](../guide/knn_search.md)。
对单查询 KNN,`GetNumElements()` 为 `1`,ids/distances 数组长度为 `k`。HGraph 和 IVF 的批量 KNN 返回
**约束:**
- 批量 IVF 搜索仅支持 KNN;不支持自定义查询距离和 reasoning labels。
- 非空外层向量必须为每个查询向量提供一个非空条目。
- 当前仅支持单查询;外层向量必须恰好包含一个条目。
Comment on lines +55 to +56
* - Batched RANGE_SEARCH is not supported; implementations MUST reject
* NumElements > 1 for range mode.
Comment thread src/algorithm/hgraph/hgraph_search.cpp
Comment thread src/algorithm/hgraph/hgraph_search.cpp Outdated
Comment on lines +893 to +895
if (search_param.time_cost != nullptr) {
search_param.time_cost->Reset();
}
Comment thread src/algorithm/ivf/ivf.cpp
Comment on lines 2141 to 2142
CHECK_ARGUMENT(request.expected_labels_.empty(),
"IVF batch search does not support expected labels");
Comment thread src/algorithm/hgraph/hgraph_search.cpp
Comment thread src/impl/label_table/label_table.cpp Outdated
Copilot AI review requested due to automatic review settings August 24, 2026 07:48
@LHT129
LHT129 force-pushed the batch_query_review_reset_1685 branch from 36a5e49 to c2b038b Compare August 24, 2026 07:48

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated 5 comments.

Suppressed comments (8)

docs/docs/en/src/api/search.md:37

  • This description says range search only accepts one query, but the IVF disable_bucket_scan routing-only mode intentionally accepts batched requests (including the RANGE_SEARCH mode) and returns bucket IDs. Clarify the exception here so this user-facing documentation does not contradict the added regression test.
| `query_` | `DatasetPtr` | `nullptr` | The query. HGraph and IVF support contiguous multi-query KNN batches; range search supports one query only. |

docs/docs/zh/src/api/search.md:37

  • 此处说明范围搜索只能接受一个查询,但 IVF 的 disable_bucket_scan 仅路由模式会有意接受批量请求(包括 RANGE_SEARCH 模式)并返回 bucket ID。请在这里说明该例外,避免用户文档与新增回归测试的行为不一致。
| `query_` | `DatasetPtr` | `nullptr` | 查询。HGraph 和 IVF 的 KNN 支持连续的多查询批次;范围搜索只支持单个查询。 |

include/vsag/search_request.h:56

  • The public contract says every batched RANGE_SEARCH must be rejected, but the IVF disable_bucket_scan routing-only path intentionally accepts batched range requests and the regression test relies on that behavior. Document this explicit exception (or reject it in the routing path) so the API contract matches the implementation.
     *          - Batched RANGE_SEARCH is not supported; implementations MUST reject
     *            NumElements > 1 for range mode.

src/algorithm/hgraph/hgraph_search.cpp:893

  • Resetting the shared timer here gives every query a fresh timeout_ms budget, so a batch can run roughly query_count * timeout_ms; additionally, ep_search_param has no timer and its routing work is not timed. This also changes single-query behavior. Preserve one request-wide deadline and pass the timer to entry-point search instead of resetting it per query.
        if (search_param.time_cost != nullptr) {
            search_param.time_cost->Reset();
        }

src/algorithm/hgraph/hgraph_search.cpp:648

  • This guard relies on HasActivePaddingLabel(), but the legacy deserialize_basic_info_v0_14() path still loads label_table_ directly without calling RebuildActivePaddingLabelIds() (unlike the modern paths). A v0.14 index containing an active external label -1 will therefore pass this check and make batch padding ambiguous. Rebuild the tracker after the legacy label vector is read before allowing batch KNN.
        CHECK_ARGUMENT(!this->label_table_->HasActivePaddingLabel(),
                       "batch KNN does not support an index containing external label -1");

src/algorithm/ivf/ivf.cpp:2193

  • Each child call creates a separate SearchStatistics/QueryContext, but the batch path discards that data and later serializes the outer stats. Batched IVF results therefore report zero distance evaluations and is_timeout=false even when a child search timed out; aggregate child statistics before returning the batch result.
            auto one_result = this->SearchWithRequest(one_request);
            if (not one_result.has_value()) {
                throw VsagException(ErrorType::INTERNAL_ERROR,
                                    "IVF batch search failed for a single query");
            }

src/algorithm/ivf/ivf.cpp:2162

  • These buffers are raw allocations and are not attached to a Dataset until after all per-query futures complete. If any nested search throws, or future.get() rethrows that failure, both ids and distances leak before the outer IndexImpl can return an error. Keep them under allocator-aware RAII (or attach ownership before starting the loop) so failed batch searches do not permanently consume query memory.
        auto* alloc = select_query_allocator(ctx.alloc, this->allocator_);
        auto* ids = static_cast<int64_t*>(alloc->Allocate(sizeof(int64_t) * total_slots));
        auto* distances = static_cast<float*>(alloc->Allocate(sizeof(float) * total_slots));

src/impl/label_table/label_table.h:454

  • This newly allocated robin set is not included in GetMemoryUsage() (which only counts deleted_ids_ around label_table.h:264-267). Index memory accounting will therefore under-report every index's label-tracking storage when active padding labels are present. Add this set's storage to the same estimate used for deleted_ids_.
    UnorderedSet<InnerIdType> deleted_ids_;  // Record deleted ids.
    UnorderedSet<InnerIdType> active_padding_label_ids_;

**Constraints:**
- Batch IVF search supports KNN only; custom query distance and reasoning labels are unsupported.
- A non-empty outer vector must contain exactly one non-empty entry per query vector.
- Currently only single-query is supported; the outer vector must contain exactly one entry.
**约束:**
- 批量 IVF 搜索仅支持 KNN;不支持自定义查询距离和 reasoning labels。
- 非空外层向量必须为每个查询向量提供一个非空条目。
- 当前仅支持单查询;外层向量必须恰好包含一个条目。
Comment thread include/vsag/index.h
Comment on lines +333 to +337
* - batched KNN requests, when supported by the implementation:
* num_elements = query->GetNumElements(),
* dim = implementation-defined returned row width. HGraph clamps it to
* min(request.topk_, GetNumElements()), while IVF preserves
* request.topk_. Callers MUST read `dim` from the returned dataset.
Comment thread src/algorithm/ivf/ivf.cpp Outdated
Comment on lines +2190 to +2194
if (not one_result.has_value()) {
throw VsagException(ErrorType::INTERNAL_ERROR,
"IVF batch search failed for a single query");
}
const auto count = std::min(request.topk_, one_result.value()->GetDim());
Comment on lines +150 to +153
for (InnerIdType id = 0; id < label_table_.size(); ++id) {
if (label_table_[id] == -1 && deleted_ids_.count(id) == 0) {
active_padding_label_ids_.insert(id);
}
Copilot AI review requested due to automatic review settings August 24, 2026 08:05
@LHT129
LHT129 force-pushed the batch_query_review_reset_1685 branch from c2b038b to 9e22ed8 Compare August 24, 2026 08:05

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated 4 comments.

Suppressed comments (11)

docs/docs/en/src/api/search.md:106

  • The IVF implementation now accepts one bucket list per query (bucket_ids_.size() == query_count), including batched requests covered by the new tests, but this constraint still says the outer vector must contain exactly one entry. This would lead clients to construct an invalid request for multi-query KNN; document one entry per query instead.
- Currently only single-query is supported; the outer vector must contain exactly one entry.

docs/docs/zh/src/api/search.md:100

  • 实现现在允许 IVF 批量请求为每个查询提供一个 bucket 列表(bucket_ids_.size() == query_count),而新增测试也覆盖了这种情况,但这里仍然写成外层向量必须恰好包含一个条目。多查询调用方会据此构造出无效请求;请改为说明每个查询对应一个条目。
- 当前仅支持单查询;外层向量必须恰好包含一个条目。

include/vsag/search_request.h:50

  • The new batch contract documents only SearchWithRequest, but the public HGraph::KnnSearch and IVF::KnnSearch overloads also delegate to this path and now accept multi-query datasets. The corresponding overload comments in include/vsag/index.h still state NumElements == 1, so update the public API documentation consistently (or explicitly exclude those overloads).
     *          - Batched KNN: Set NumElements to the number of queries, with vectors
     *            stored contiguously. Supported by HGraph::SearchWithRequest and
     *            IVF::SearchWithRequest; results are returned with NumElements =
     *            query_count and a row-major Dim determined by the implementation

include/vsag/search_request.h:55

  • Both HGraph and IVF reject expected_labels_ for batched KNN, but this batch contract does not state that restriction. A caller following the documented batch behavior and enabling reasoning receives an invalid-argument error; document that expected-label reasoning (and custom-distance callbacks, which are also single-query-only here) requires a single query.
     *            (id = -1, distance = +infinity). Batch KNN rejects an index containing external
      *            label -1 to keep this padding unambiguous.
     *          - Batched RANGE_SEARCH is not supported; implementations MUST reject

src/algorithm/hgraph/hgraph_search.cpp:893

  • base_search_param.time_cost starts before the routing loop, but resetting it here discards routing time for every request, including single-query requests. A timeout can therefore be exceeded during routing and still receive a fresh full timeout_ms budget for approximate search, changing the previous single-query deadline behavior. Start/reset the timer before each query's routing (or otherwise account for routing) while preserving the intended batch budget.
            search_param.time_cost->Reset();
        }

src/algorithm/hgraph/hgraph_serialize.cpp:320

  • This rebuild only covers the footer-based deserialization paths. The legacy deserialize_basic_info_v0_14 path still reads label_table_ directly and never rebuilds active_padding_label_ids_; a v0.14 index containing a live external label -1 will therefore pass HasActivePaddingLabel() and batch KNN can emit an ambiguous -1 value. Rebuild the active-label set immediately after the legacy label-table read as well.
        this->label_table_->RebuildActivePaddingLabelIds();

src/algorithm/ivf/ivf.cpp:2189

  • Each batched row is executed through a recursive SearchWithRequest, which creates a fresh local SearchStatistics; the outer stats object dumped at line 2219 is never updated. Consequently a successful batched IVF result reports zero distance evaluations and other counters, and does not expose per-query timeout state, even though the nested searches performed work. Aggregate the nested statistics into the batch context (or route each search through the outer QueryContext) before returning.
            auto one_result = this->SearchWithRequest(one_request);

src/impl/label_table/label_table.h:454

  • The newly allocated active_padding_label_ids_ set is not included in LabelTable::GetMemoryUsage(), which currently accounts for deleted_ids_ but not this set's dynamic storage. Index memory is therefore underreported whenever active external -1 labels are present. Include this set's allocated footprint in the memory-usage calculation.
    UnorderedSet<InnerIdType> active_padding_label_ids_;

src/utils/timer.cpp:52

  • Reset() is newly added and is relied on by batched HGraph timeout handling, but src/utils/timer_test.cpp only covers construction, Record(), SetThreshold(), and destructor recording. Add a regression test that verifies a timeout threshold is measured from the reset point; otherwise the per-query timeout behavior can regress without a failing test.
Timer::Reset() {

tests/test_hgraph.cpp:4355

  • dim is the loop variable for the current fixture dimension, but this assignment mutates it for the remainder of the test_cases loop. Since the RaBitQ cases occur before pq and the other quantizers, those later cases run at 960 instead of the requested fixture dimension, silently losing coverage. Compute a separate effective dimension for the RaBitQ case and leave dim unchanged.
                if (HGraphTestIndex::IsRaBitQ(base_quantization_str) &&
                    dim < fixtures::RABITQ_MIN_RACALL_DIM) {
                    dim = fixtures::RABITQ_MIN_RACALL_DIM;

tests/test_hgraph.cpp:4438

  • This second new test repeats the same loop-variable mutation: once a RaBitQ case raises dim to 960, subsequent quantization cases in this dimension iteration no longer use the fixture dimension. That makes the test pass while skipping intended parameter coverage. Use a separate effective dimension for the RaBitQ build and dataset.
                if (HGraphTestIndex::IsRaBitQ(base_quantization_str) &&
                    dim < fixtures::RABITQ_MIN_RACALL_DIM) {
                    dim = fixtures::RABITQ_MIN_RACALL_DIM;

Comment on lines +645 to +648
CHECK_ARGUMENT(request.expected_labels_.empty(),
"reasoning (expected_labels_) is only supported for single-query search");
CHECK_ARGUMENT(!this->label_table_->HasActivePaddingLabel(),
"batch KNN does not support an index containing external label -1");
Comment thread src/algorithm/ivf/ivf.cpp
Comment on lines 2161 to 2162
auto* ids = static_cast<int64_t*>(alloc->Allocate(sizeof(int64_t) * total_slots));
auto* distances = static_cast<float*>(alloc->Allocate(sizeof(float) * total_slots));
Comment thread src/algorithm/ivf/ivf.cpp
@@ -2177,6 +2187,7 @@ IVF::SearchWithRequest(const SearchRequest& request) const {
}
one_request.params_str_ = json.Dump();
auto one_result = this->SearchWithRequest(one_request);
Comment thread src/algorithm/ivf/ivf.cpp
Comment on lines +2144 to +2145
CHECK_ARGUMENT(!this->label_table_->HasActivePaddingLabel(),
"batch KNN does not support an index containing external label -1");
Comment thread src/algorithm/ivf/ivf.cpp
}
one_request.params_str_ = json.Dump();
auto one_result = this->SearchWithRequest(one_request);
CHECK_ARGUMENT(one_result != nullptr, "IVF batch search returned an empty result");

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] The IVF batch path recursively calls this->SearchWithRequest(one_request) for each query. Each recursive call goes through the full SearchWithRequest including argument validation, create_search_param, and the single-query search path. For empty indexes, this means num_queries full traversals of the validation/ dispatch logic before each per-query search() returns an empty result. The HGraph batch path handles empty indexes more efficiently by returning a pre-shaped empty result early. Consider adding an early-return for GetNumElements() == 0 before entering the batch loop to avoid the redundant per-query overhead.

Comment thread src/algorithm/ivf/ivf.cpp
one_request.params_str_ = json.Dump();
auto one_result = this->SearchWithRequest(one_request);
CHECK_ARGUMENT(one_result != nullptr, "IVF batch search returned an empty result");
const auto count = std::min(request.topk_, one_result->GetDim());

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] The error message says "returned an empty result" but the check is one_result != nullptr, which detects a null pointer, not an empty dataset. Consider changing the message to "IVF batch search returned null" or "IVF batch search per-query result is null" for accuracy.

search_param.enable_rabitq_one_bit_search = params.rabitq_one_bit_search;
if (params.enable_time_record) {
search_param.time_cost = std::make_shared<Timer>();
search_param.time_cost->SetThreshold(params.timeout_ms);

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] The visited_list_guard struct is defined identically in both search_range_with_request (line 456) and SearchWithRequest (line 833). Consider extracting it to a shared location (e.g., as a private nested struct in hgraph.h or a file-scope helper) to avoid duplication and ensure consistent RAII behavior across all search paths.

Comment thread tests/test_hgraph.cpp
HGraphTestIndex::TestGeneral(cache_index, dataset, search_param, 0.98f);
}

static void

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] The TestHGraphMultiQueryKnnSearch test exercises multi-query via the KnnSearch overload, but the primary multi-query entry point documented in the API is SearchWithRequest. Consider adding a test case that uses SearchWithRequest with NumElements > 1 to validate the full public API path, including the SearchMode::KNN_SEARCH mode and the rectangular result layout contract (row-major, padding with id=-1).

Signed-off-by: LHT129 <tianlan.lht@antgroup.com>

Co-authored-by: opencode <opencode@anthropic.com>
Signed-off-by: LHT129 <tianlan.lht@antgroup.com>
Copilot AI review requested due to automatic review settings August 25, 2026 09:32
@LHT129
LHT129 force-pushed the batch_query_review_reset_1685 branch from 9e22ed8 to 0906b18 Compare August 25, 2026 09:32
Comment thread src/algorithm/ivf/ivf.cpp
@@ -2141,12 +2141,22 @@ IVF::SearchWithRequest(const SearchRequest& request) const {
CHECK_ARGUMENT(request.expected_labels_.empty(),

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] JsonType::Parse + json.Dump() inside the per-query lambda parses and re-serializes the params string on every iteration. For large batches this adds measurable overhead.

The only mutation is setting parallelism to 1. Consider hoisting the JSON mutation outside the loop — parse once, mutate, dump once, and reuse the resulting string across all queries.

const FilterPtr& filter,
QueryContext& ctx) const {
InnerSearchParam search_param;
search_param.ep = this->entry_point_id_;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[note] In search_range_with_request, enable_reorder and enable_rabitq_one_bit_search are set directly from params without the use_custom_distance guard (lines 59-60 in this hunk). While this is currently unreachable for custom distance because range+custom_distance is rejected at line 582-583 in SearchWithRequest, it creates a latent inconsistency with the KNN batch path which properly guards these fields (lines 645-647: use_custom_distance ? false : params.enable_reorder). Consider adding the guard for consistency and future-proofing.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

kind/feature Brand-new functionality or capabilities 引入全新的功能、新特性或新能力 module/api module/docs module/index module/testing size/XXL version/1.1

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature] Add multi-query batch search support

3 participants